"""
对简单的仿射模型进行 随机 梯度下降算法,求出最佳权重、单个样本最小损失loss
f(x)=x*w
loss=(f(xn)-y)**2
"""
import random

import numpy as np
from matplotlib import pyplot as plt

# 定义已知的数据集,training的dataset
x_data = [1.0, 2.0, 3.0]
y_data = [2.0, 4.0, 6.0]

# 定义初始的权重w
w = 1.0


# 定义前馈函数
def forward(x):
    return x * w


# 定义损失函数
def loss(x, y):
    """
    对当前的权重w,进行单样本的误差计算
    :param x: x值
    :param y: y值
    :return: 单样本的误差
    """
    # 计算当前权重下的单个样本的前馈
    y_hat = forward(x)
    # 计算单个样本的loss损失并记录
    tempLoss = (y_hat - y) ** 2
    return tempLoss


# 定义梯度函数(求取loss对w的导)
def gradient(x, y):
    theGradient = 2 * x * (w * x - y)
    # 返回当前的梯度
    return theGradient


allCost = []
epochs = []
print("predict(before training): ", forward(4))
# 进行循环训练
for epoch in range(100):
    # 选择一个随机的x,y
    idx = random.randrange(0, len(x_data))
    x = x_data[idx]
    y = y_data[idx]

    # 计算一个样本的梯度
    tempGradient = gradient(x, y)
    # 0.01是学习率
    w -= 0.01 * tempGradient

    # 计算当前的损失
    tempCost = loss(x, y)
    # 记录值
    epochs.append(epoch)
    allCost.append(tempCost)
    print(f"Epoch: {epoch}, w: {w}, Cost: {tempCost}")
print("predict(after training): ", forward(4))

# 绘制损失-轮数图
plt.plot(epochs, allCost)
plt.ylabel("cost")
plt.xlabel("epoch")
plt.grid()
plt.show()
